from pymilvus import MilvusClient

milvus_username = 'root'
milvus_password = 'Milvus'
client = MilvusClient(
    uri="http://localhost:19530",
    token=f"{milvus_username}:{milvus_password}",
)

def query_str(query_str):
    from pymilvus.model.hybrid import BGEM3EmbeddingFunction

    bge_m3_ef = BGEM3EmbeddingFunction(
        # model_name='BAAI/bge-m3', # Specify the model name
        model_name='BAAI/bge-base-en-v1.5', # Specify the model name
        device='cpu', # Specify the device to use, e.g., 'cpu' or 'cuda:0'
        use_fp16=False # Specify whether to use fp16. Set to `False` if `device` is `cpu`.
    )
    queries = [query_str]
    query_embeddings = bge_m3_ef.encode_queries(queries)
    return query_embeddings

def test01(query):
    milvus_collection_name = 'demo_collection'
    query_vector=query_str(query)["dense"]
    print(query_vector)
    res = client.search(
        collection_name=milvus_collection_name, # Collection name
        data=query_vector, # Replace with your query vector
        search_params={
            "metric_type": "COSINE",
            "params": {"level":1}, # Search parameters
        }, # Search parameters
        limit=15, # Max. number of search results to return
        output_fields=["id","text"], # Fields to return in the search results
        consistency_level="Bounded"
    )
    return res

def test02(query):
    milvus_collection_name = 'advance_rag'
    # query_vector=query_str(query)["dense"]
    query_vector=query_str(query)
    print(query_vector)
    # python search code
    res = client.search(
        collection_name=milvus_collection_name,  # Collection name
        data=query_vector,  # Replace with your query vector
        search_params={
            "metric_type": "IP",
            "params": {"nprobe": 1},  # Search parameters
        },  # Search parameters
        limit=15,  # Max. number of search results to return
        output_fields=["id"],  # Fields to return in the search results
        consistency_level="Bounded"
    )

    return res


if __name__ == '__main__':
    query="When was artificial intelligence founded"
    # res=test01(query)
    res=test02(query)
    print(res)

